import numpy as np
import tensorflow as tf
from tensorflow.keras.datasets import mnist
from tensorflow.keras.layers import Flatten, Dense
from tensorflow.keras.models import Sequential
from tensorflow.keras.utils import to_categorical

# 1. 引入数据，对数据进行预处理
(X_train, y_train), (X_test, y_test) = mnist.load_data()
X_train = X_train / 255.0
X_test = X_test / 255.0
y_train = to_categorical(y_train)
y_test = to_categorical(y_test)

# 2. 构建网络模型
model = Sequential()
model.add(Flatten(input_shape=(28, 28)))
model.add(Dense(20, activation=tf.nn.relu))
model.add(Dense(10, activation=tf.nn.softmax))

# 3. 编译模型，指定优化器、损失函数和评估指标
model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

# 4. 训练模型
model.fit(X_train, y_train, epochs=20, validation_split=0.2)

# 5. 评估模型
test_loss, test_acc = model.evaluate(X_test, y_test)
print('Test accuracy:', test_acc)

# 6. 保存模型到硬盘
model_path = './mnist_model'
tf.keras.models.save_model(model, model_path)
